#include <iostream>
#include <string>

#include "Algorithm.hpp"
#include "RSA.hpp"
#include "Message.hpp"
#include "OperationMode_ECB.hpp"
#include "OperationMode_CBC.hpp"
#include "OperationMode_CFB.hpp"
#include "OperationMode_OFB.hpp"
#include "AES.hpp"

enum CipherMode { CIPHER_AES, CIPHER_RSA };
enum ProgramMode { SETUP, ENCRYPT, DECRYPT };
enum OpMode { NONE, ECB, CBC, CFB, OFB };

bool checkParameters(int argc, char **argv) {
	if ((argc != 5) && (argc != 7)) {
		std::cerr << "Invalid parameter count." << std::endl;
		return false;
	}
	if ((argv[1] != std::string("RSA")) && (argv[1] != std::string("AES"))) {
		std::cerr << "System type must be either 'RSA' or 'AES'." << std::endl;
		return false;
	}

	if ((argv[2] != std::string("-s")) && (argv[2] != std::string("-e")) && (argv[2] != std::string("-d"))) {
		std::cerr << "Mode must be either setup, encrypt or decrypt." << std::endl;
		return false;
	}

	if ((argv[1] == std::string("AES")) && (argv[2] == std::string("-s"))) {
		std::cerr << "AES does not support key setup." << std::endl;
		return false;
	}

	if ((argv[2] == std::string("-s")) && (argc != 5)) {
		std::cerr << "Invalid parameter count for key setup." << std::endl;
		return false;
	}

	if ((argv[1] == std::string("AES")) && (std::string(argv[3]).size() != 16)) {
		std::cerr << "AES128 requires password length of 16 characters." << std::endl;
		return false;
	}

	if ((argv[2] == std::string("-e")) || (argv[2] == std::string("-d"))) {
		if (argc != 7) {
			std::cerr << "Invalid parameter count for encrypt/decrypt." << std::endl;
			return false;
		}
		if ((std::string("ecb") != argv[4]) 
			&& (std::string("cbc") != argv[4])
			&& (std::string(argv[4]).substr(0, 4) != "cfb-")
			&& (std::string(argv[4]).substr(0, 4) != "ofb-")) {
			std::cerr << "Operation mode must be one of ecb, cbc, cfb or ofb." << std::endl;
			return false;
		}
	}

	return true;
}

int protected_main(int argc, char **argv) {
	if (!checkParameters(argc, argv)) {
		std::cerr << argv[0] << " RSA [-s] [Bitlength] [RSA Parameters Output File]" << std::endl;
		std::cerr << argv[0] << " RSA [-e] [RSA Parameters Input File] [ecb|cbc|cfb-j|ofb-j] [Plain Filename] [Cipher Filename]" << std::endl;
		std::cerr << argv[0] << " RSA [-d] [RSA Parameters Input File] [ecb|cbc|cfb-j|ofb-j] [Cipher Filename] [Plain Filename]" << std::endl;
		std::cerr << argv[0] << " AES [-e] [AES Key] [ecb|cbc|cfb-j|ofb-j] [Plain Filename] [Cipher Filename]" << std::endl;
		std::cerr << argv[0] << " AES [-d] [AES Key] [ecb|cbc|cfb-j|ofb-j] [Cipher Filename] [Plain Filename]" << std::endl;
		std::cerr << std::endl;
		std::cerr << "e.g.:" << std::endl;
		std::cerr << "    " << argv[0] << " RSA -s 1024 keyring.txt" << std::endl;
		std::cerr << "    " << argv[0] << " RSA -e keyring.txt cbc secret.txt secret.enc" << std::endl;
		std::cerr << "    " << argv[0] << " RSA -d keyring.txt cbc secret.enc secret_decoded.txt" << std::endl;
		std::cerr << "    " << argv[0] << " RSA -e keyring.txt ofb-32 secret2.txt secret2.enc" << std::endl;
		std::cerr << "    " << argv[0] << " RSA -d keyring.txt ofb-32 secret2.enc secret2_decoded.txt" << std::endl;
		std::cerr << "e.g.:" << std::endl;
		std::cerr << "    " << argv[0] << " AES -e foobarpassword99 cbc secret2.txt secret2.enc" << std::endl;
		std::cerr << "    " << argv[0] << " AES -d foobarpassword99 cbc secret2.enc secret2_decoded.txt" << std::endl;
		exit(EXIT_FAILURE);
	}

	ProgramMode mode;
	if (std::string("-s") == argv[2]) {
		/* Setup RSA cryptosystem */
		std::string parameters(argv[4]);
		unsigned int bitlength = atoi(argv[3]);
		RSA RSA_System = RSA();
		std::cerr << "Generating cryptosystem of length " << bitlength << " bits..." << std::endl;
		RSA_System.setup(bitlength);
		RSA_System.writeParametersToFile(parameters);
		return 0;
	} else if (std::string("-e") == argv[2]) {
		mode = ENCRYPT;
	} else {
		mode = DECRYPT;
	}

	CipherMode cipher;
	if (std::string("AES") == argv[1]) {
		cipher = CIPHER_AES;
	} else {
		cipher = CIPHER_RSA;
	}
	
	unsigned int blockSizeBits = 128;
	OpMode opMode;
	if (std::string("ecb") == argv[4]) opMode = ECB;
		else if (std::string("cbc") == argv[4]) opMode = CBC;
		else if (std::string(argv[4]).substr(0, 4) == "cfb-") {
			opMode = CFB;
			blockSizeBits = atoi((std::string(argv[4])).substr(4).c_str());
		} else {
			opMode = OFB;
			blockSizeBits = atoi((std::string(argv[4])).substr(4).c_str());
		}
	
	if ((blockSizeBits % 8) != 0) {
		std::cerr << "j must be a multiple of 8 bits, " << blockSizeBits << "is not." << std::endl;
	}

	Algorithm *systemPtr;
	if (cipher == CIPHER_AES) {
		systemPtr = new AES();
		AES &aesSystem = *(dynamic_cast<AES*>(systemPtr));
		aesSystem.setKey(Block(argv[3]));
	} else {
		systemPtr = new RSA();
		RSA &rsaSystem = *(dynamic_cast<RSA*>(systemPtr));
		rsaSystem.readParametersFromFile(argv[3]);
	}
	Algorithm &system = *systemPtr;

	OperationMode *om = NULL;
	switch (opMode) {
		case NONE:	om = NULL; break;
		case ECB:	om = new OperationMode_ECB(system); break;
		case CBC:	om = new OperationMode_CBC(system); break;
		case CFB:	om = new OperationMode_CFB(system, blockSizeBits); break;
		case OFB:	om = new OperationMode_OFB(system, blockSizeBits); break;
	}

	std::string infile = argv[5];
	std::string outfile = argv[6];
	if (mode == ENCRYPT) {
		Message msg(blockSizeBits / 8);
		msg.readFromFile(infile);
		std::cerr << "Encrypting " << infile << ": " << msg.getBlockCount() << " blocks, " << msg.getSize() << " bytes" << std::endl;
	
		FILE *f;
		f = fopen(outfile.c_str(), "w");
		if (f == NULL) return 1;

		fprintf(f, "%d\n", msg.getSize());
		for (unsigned int i = 0; i < msg.getBlockCount(); i++) {
			Block plainBlock;
			plainBlock = msg.getBlock(i);
	
			Block cipherBlock = om->encrypt(plainBlock);

			Integer cipherInt(cipherBlock);
			cipherInt.writeToFile(f);
		}

		fclose(f);
	} else if (mode == DECRYPT) {
		Message msg(blockSizeBits / 8);
		FILE *f;
		f = fopen(infile.c_str(), "r");
		if (f == NULL) return 1;

		unsigned int msgSize;
		if (fscanf(f, "%d\n", &msgSize) != 1) throw GENERICEXCEPTION("parser error");
		msg.prepare(msgSize);

		for (unsigned int i = 0; i < msg.getBlockCount(); i++) {
			Integer cipherInt;
			cipherInt.readFromFile(f);

			Block cipherBlock;
			cipherBlock = cipherInt.toBlock();
			Block plainBlock = om->decrypt(cipherBlock);
			msg.setBlock(i, plainBlock);
		}
		msg.writeToFile(outfile);
		fclose(f);

		std::cerr << "Decrypted " << infile << ": " << msg.getBlockCount() << " blocks, " << msg.getSize() << " bytes" << std::endl;
	}
	return 0;
}

int main(int argc, char **argv) {
//	return protected_main(argc, argv);
	
	try {
		return protected_main(argc, argv);
	} catch (const GenericException &E) {
		E.terminate();
	}
	return 1;
}
